import torch

from opt_func.FuncSet_torch import *


class Function :
    def __init__(self,
                 function,
                 values_range,
                 data_index=None,
                 restrains=None,
                 restrains_all_rebase_score=1e8) :
        """
        优化函数对象的初始化定义，包括：
        1.优化函数
        2.各变量取值空间
        3.变量次序（用于多树训练）
        4.约束列表（元素都是约束函数）
        5.违反约束时的惩罚数值
        """
        function_list = ['Rastrigin','MishraBird','Goldsteitorchrice','Rosenbrock']
        if isinstance(function,str) :
            if function not in function_list :
                raise AttributeError('输入的函数不存在')
            else :
                self.function = eval(function)
        else :
            self.function = function
        if data_index is None:
            raise ValueError("请制定变量序列")
        self.restrains_all_rebase_score = restrains_all_rebase_score
        self.values_range = values_range #用以告诉树变量的数量
        self.data_index = data_index #变量的顺序

        #todo:加入约束的作用
        self.restrains = restrains


    def __call__(self,fixed,simulation=None,**kwargs) :
        """:
        将对象作为函数进行调用
        在该方法中，实现了：
        1.生成优化函数传入数组（这里的实现不怎么优美，可以看看怎么优化）
        2.当有约束时考虑约束的影响
        """
        # TODO:对于违反程度做个判定方法
        # TODO:如果所有都违反约束怎么办
        fixed = torch.tensor(fixed)
        data = fixed
        #如果进行了蒙特卡洛仿真，则对数组进行堆叠构造
        if simulation is not None:
            _,wide = simulation.shape #wide是仿真的次数
            fixed = fixed.repeat(wide,1).t()
            data = torch.vstack([fixed,simulation])
            print(data)
            print(fixed)
        data = data[self.data_index]
        score = self.function(data,**kwargs)
        #约束的处理
        score = self.process_restrain(data,score,**kwargs)
        return [score.min(),score.float().mean()]

    def process_restrain(self,data,score,**kwargs):
        if isinstance(self.restrains,dict) :  # 以字典形式反映惩罚系数
            is_rebase_restrains = torch.tensor([restrain(data,**kwargs)
                                            for restrain in self.restrains.keys()])*torch.array(
                self.restrains.values())
            score += is_rebase_restrains.sum(axis=0)
        elif isinstance(self.restrains,list) :  # 没有惩罚系数
            is_rebase_restrains = torch.tensor([restrain(data,**kwargs)
                                            for restrain in self.restrains]).any(axis=0)
            score = score[is_rebase_restrains]
        #TODO:决定全部违反时怎么办
            if is_rebase_restrains.sum() == 0 :
                score = self.restrains_all_rebase_score
        return score
if __name__ == '__main__':
    # function,
    # values_range,
    # data_index = None,
    # restrains = None,
    # restrains_all_rebase_score = 1e8):
    f = Rosenbrock
    data_index = range(5)
    test = Function(f,None,data_index)
    print(test(range(3),simulation=torch.arange(8).view(2,4)))
